# 功能代码
import cv2   
import mediapipe as mp   
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.animation import FuncAnimation
from PIL import Image
import datatreating as dt
from compare import set_compare
from compare import optical_flow_keyframes as keyframes
import time
from matplotlib.backends.backend_agg import FigureCanvasAgg
import base64
from io import BytesIO
import datetime
import os
import csv


human_height_pixels = 0   # 人体像素身高 全局变量
height_real = 170         # 人体实际身高 全局变量
now_ankle_pixel_coord = (0,0)
landmark_num_distance = 0

side = 0
# 攻击时的数据
attacks_count = 0
global_sum_distance = 0
global_right_hand_angle = 0
global_right_ankle_angle = 0
global_left_ankle_angle = 0
global_right_knee_angle = 0
global_left_knee_angle = 0
global_right_shoulder = (0,0)
global_left_shoulder = (0,0)
global_right_hip = (0,0)
global_left_hip = (0,0)

# 定义一个姿态检测的类
class PoseDetector:


    # 初始化函数，video_path是视频路径
    def __init__(self, video_path,height_input,side_input):
        global height_real
        global side
        height_real = height_input
        side = side_input
        self.mp_pose = mp.solutions.pose   # 使用pose方法（识别姿势）
        self.mp_draw = mp.solutions.drawing_utils  #画图方法
        # pose的初始化设置
        # min_detection_confidence = 人员检测模型的最小置信值
        # min_tracking_confidence = 地标跟踪模型的最小置信值
        self.pose = self.mp_pose.Pose(
            min_detection_confidence=0.5,
            min_tracking_confidence=0.5
        )
        
        self.height_real = height_real

        # 读取视频
        self.cap = cv2.VideoCapture(video_path)

        # 获取视频的原始尺寸
        self.frame_width = int(self.cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        self.frame_height = int(self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

        # 获取视频总帧数
        self.total_frames = int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT)) 
        
        # 初始化matplotlib图形
        self.fig = plt.figure(figsize=(10, 10))
        self.ax = self.fig.add_subplot(111, projection='3d')
        
        #存储关键点数据
        self.landmarks_3d = []
        
        # 存储视角参数
        self.elev = 15
        self.azim = 90
        
        # 创建OpenCV窗口并设置原始尺寸
        # cv2.namedWindow('Pose Detection', cv2.WINDOW_NORMAL)
        # cv2.resizeWindow('Pose Detection', 640, 480)
        
        # 连接matplotlib的视角更新事件
        self.fig.canvas.mpl_connect('motion_notify_event', self._update_view)

        # 初始化CSV文件
        # self.csv_W_AllLandmarks,self.csv_W_NeedData,self.csv_W_AttackData = dt.setup_csv(self)
        self.csv = dt.setup_csv()
        # 存放当前帧的归一化坐标
        self.current_landmarks = []


    # 旋转3D视图
    def _update_view(self, event):
        if event.inaxes == self.ax:
            self.elev = self.ax.elev
            self.azim = self.ax.azim


    # 处理当前帧
    def process_frame(self):

        global human_height_pixels
        global now_ankle_pixel_coord
        global landmark_num_distance
        global attacks_count
        
        global global_sum_distance
        global global_right_hand_angle
        global global_right_ankle_angle
        global global_left_ankle_angle
        global global_right_knee_angle
        global global_left_knee_angle
        global global_right_shoulder
        global global_left_shoulder
        global global_right_hip
        global global_left_hip


        success, image = self.cap.read()  #读取当前帧
        if not success:
            return False
        
        self.rezise_image = cv2.resize(image,(640,480))

        # 根据变量side切割图像
        image11 = dt.slicing_photo(self,side,self.rezise_image)

        # 转换颜色空间用于MediaPipe处理 将BGR格式转换为RGB格式 这是因为 MediaPipe 库需要 RGB 格式的图像作为输入。
        image_rgb = cv2.cvtColor(image11, cv2.COLOR_BGR2RGB)
        

        # 进行姿态检测 处理rgb图像，并返回检测结果
        self.results = self.pose.process(image_rgb)
        
        # 在图像上绘制检测结果 
        if self.results.pose_landmarks:
            self.mp_draw.draw_landmarks(
                image11, 
                self.results.pose_landmarks, 
                self.mp_pose.POSE_CONNECTIONS
            )
            
            # 显示图像
            # cv2.imshow('Pose Detection', image11)
            # cv2.waitKey(1) 
                                   
            if self.results.pose_world_landmarks:
                landmarks = self.results.pose_world_landmarks.landmark

                # 获取当前帧数 将帧数int类型输出
                self.frame_number = int(self.cap.get(cv2.CAP_PROP_POS_FRAMES))

                # 提取所有关键点的3D坐标并进行坐标变换
                x = [-landmark.x for landmark in landmarks]
                y = [-landmark.y for landmark in landmarks]
                z = [landmark.z for landmark in landmarks]
                
                # 交换y和z坐标，使人体垂直显示
                self.landmarks_3d = [x,z,y]
                # print(f"111:{self.landmarks_3d}")

#############################################################################################################################################
            
                # 获取第一帧，用来计算 像素身高 和 比例因子
                if human_height_pixels == 0:
                    print("yes")
                    self.pixel_ratio,human_height_pixels = dt.hight_pixel_rstio(self) # 返回值为 比例因子
                    self.pixel_ratio = float(self.pixel_ratio)


                #     # # 判断人脸朝向 并确定前脚
                #     nose = dt.normalized_to_pixel(self,0)[0]
                #     right_ear = dt.normalized_to_pixel(self,8)[0]
                #     right_ankle = dt.normalized_to_pixel(self,28)[0]
                #     left_ankle = dt.normalized_to_pixel(self,27)[0]
                #     if nose[0] < right_ear[0]:
                #         self.face = "right"
                #         if right_ankle[0] < left_ankle[0]:
                #             self.forefoot = 32     # 人在右边时，右脚在前
                #         else:
                #             self.forefoot = 31     # 人在右边时，左脚在前
                #     else:
                #         self.face = "left"
                #         if right_ankle[0] > left_ankle[0]:
                #             self.forefoot = 32     # 人在左边时，右脚在前
                #         else:
                #             self.forefoot = 31     # 人在左边时，左脚在前
                #     print(self.face)


                    self.face_state = 0
                    left_heel = dt.normalized_to_pixel(self,29)[0]
                    right_heel = dt.normalized_to_pixel(self,30)[0]
                    left_foot_index = dt.normalized_to_pixel(self,31)[0]
                    right_foot_index = dt.normalized_to_pixel(self,32)[0]

                    # 人物在左边
                    if (left_foot_index > left_heel or  right_foot_index > right_heel):
                        if left_heel > right_heel:
                            self.face_state = 1  # 左左撇子
                            self.forefoot = 31
                        else:
                            self.face_state = 2  # 左右撇子
                            self.forefoot = 32     # 人在左边时，右脚在前

                    # 人物在右边
                    if(left_foot_index < left_heel or right_foot_index < right_heel):
                        if left_heel < right_heel:
                            self.face_state = 3  # 右左撇子
                            self.forefoot = 31

                        else:
                            self.face_state = 4  # 右右撇子
                            self.forefoot = 32

                    print(self.face_state)


                # 计算所需角度
                right_ankle_angle = dt.cal_ang(self,26,28,-1)   # 计算右脚踝与地面夹角
                left_ankle_angle = dt.cal_ang(self,25,27,-1)    # 计算左脚踝与地面夹角
                right_knee_angle = dt.cal_ang(self,24,26,28)   # 计算右膝角度
                left_knee_angle = dt.cal_ang(self,23,25,27)     # 计算左膝角度
                right_hand_angle = dt.cal_ang(self,12,14,16)    # 计算右肘角度

                # # 获取躯干四点
                right_shoulder = dt.normalized_to_pixel(self,12)[0]  #获取右肩坐标
                left_shoulder = dt.normalized_to_pixel(self,11)[0]  #获取左肩坐标
                right_hip = dt.normalized_to_pixel(self,24)[0]      #获取右臀坐标
                left_hip = dt.normalized_to_pixel(self,24)[0]       #获取左臀坐标
                self.csv[4].writerow([self.frame_number,
                                            left_shoulder,
                                            right_shoulder,
                                            left_hip,
                                            right_hip,
                                            left_knee_angle,
                                            right_knee_angle,
                                            right_hand_angle])
               
                for landmark_value in range(33):
                    # 获取当前关键点的像素坐标和归一化坐标
                    landmark_pixel_coord,normalized_coord = (dt.normalized_to_pixel(
                                                             self,landmark_value))

                    # 将归一化坐标写入列表
                    self.current_landmarks.append([float(normalized_coord[0]),
                                                   float(normalized_coord[1])])

                    # 只能开其中一个
                    # 将像素坐标写入AllLandmarks.csv   
                    self.csv[3].writerow([self.frame_number,
                                          landmark_value,
                                          landmark_pixel_coord[0],
                                          landmark_pixel_coord[1]])
                    # 将归一化坐标写入AllLandmarks.csv
                    # self.csv[3].writerow([frame_number,landmark_value,normalized_coord[0],normalized_coord[1]])

                    if self.forefoot == landmark_value:
                         # 更新前脚坐标点
                        forefoot_pixel_coord = dt.normalized_to_pixel(self,self.forefoot)[0]
                        last_ankle_pixel_coord = now_ankle_pixel_coord
                        now_ankle_pixel_coord = forefoot_pixel_coord

                        # 计算移动距离
                        # 判断向前移动
                        # print(111)
                        if (now_ankle_pixel_coord[0] <= last_ankle_pixel_coord[0] and 
                            self.face_state in [3, 4]):  # 向左移动
                        # if now_ankle_pixel_coord[0] <= last_ankle_pixel_coord[0] and self.face == "right":  # 向左移动
                        
                            # 累加每一帧距离
                            landmark_num_distance += (dt.Two_point_distance(
                                                      now_ankle_pixel_coord,last_ankle_pixel_coord)*
                                                      self.pixel_ratio)
                        
                        elif (now_ankle_pixel_coord[0] >= last_ankle_pixel_coord[0] and 
                              self.face_state in [1,2]):  # 向右移动
                        # elif now_ankle_pixel_coord[0] >= last_ankle_pixel_coord[0] and self.face =="left":  # 向右移动
                            
                            landmark_num_distance += (dt.Two_point_distance(
                                                      now_ankle_pixel_coord,last_ankle_pixel_coord)*
                                                      self.pixel_ratio)
                        
                        # 移动结束，进入结算
                        
                        else:
                            if landmark_num_distance >=20 and right_hand_angle >= 165:
                            # if right_hand_angle >= 165:
                                # 累加攻击次数
                                attacks_count += 1
                                sum_distance = format(landmark_num_distance,'.2f')  # 像素移动距离*比例因子

                                print(f"第{attacks_count}次进攻：")
                                print(f"移动距离{sum_distance}\t右手夹角{right_hand_angle}°")
                                print(f"左脚与地面夹角{left_ankle_angle}°\t右脚与地面夹角{right_ankle_angle}°")
                                print(f"左膝盖夹角{left_knee_angle}°\t右膝盖夹角{right_knee_angle}°")
                                print(f"左肩膀坐标{left_shoulder}\t右肩坐标{right_shoulder}")
                                print(f"左臀部坐标{left_hip}\t右臀部坐标{right_hip}")
                                # print(222)
#                                 print(f'''第{attacks_count}次进攻：
# 移动距离{sum_distance} 
# 右脚夹角{right_ankle_angle}° 左脚夹角{left_ankle_angle}° 右膝盖夹角{right_knee_angle}° 
# 左膝盖夹角{left_knee_angle}° 右手夹角{right_hand_angle}° 
# 右肩坐标{right_shoulder} 左肩坐标{left_shoulder} 右臀坐标{right_hip} 左臀坐标{left_hip}
# ''' ) 
                                
                                global_sum_distance = sum_distance
                                global_right_hand_angle = right_hand_angle
                                global_right_ankle_angle = right_ankle_angle
                                global_left_ankle_angle = left_ankle_angle
                                global_right_knee_angle = left_ankle_angle
                                global_left_knee_angle = left_knee_angle
                                global_right_shoulder = right_shoulder
                                global_left_shoulder = left_shoulder
                                global_right_hip = right_hip
                                global_left_hip = left_hip
                                # 将攻击数据写入csv
                                # self.csv_W_AttackData.writerow([attacks_count,
                                #                                 right_hand_angle,
                                #                                 right_ankle_angle,
                                #                                 left_ankle_angle,
                                #                                 sum_distance])
                                self.csv[5].writerow([attacks_count,
                                                      right_hand_angle,
                                                      right_ankle_angle,
                                                      left_ankle_angle,
                                                      sum_distance])
                            # 移动结束，距离清零
                            landmark_num_distance = 0

                # 将归一化坐标列表转换成ndarrays数组
                # print(333)
                NPcurrent_landmarks = np.array(self.current_landmarks)   
                # 判断动作标准度
                compare_result = set_compare(self.face_state,NPcurrent_landmarks)
                # print(444)
                # print(compare_result)
                # set_compare(self.face,NPcurrent_landmarks)
                if compare_result[0] != 0:
                    if compare_result[0] == 3:
                        print(f"准备动作标准，标准度{compare_result[1]}")
                        print(f"攻击动作标准，标准度{compare_result[2]}")
                    if compare_result[0] == 1:
                        print(f"准备动作标准，标准度{compare_result[1]}")
                    if compare_result[0] == 2:
                        print(f"攻击动作标准，标准度{compare_result[2]}")

                # 归一化坐标列表清空
                self.current_landmarks.clear() 
 #############################################################################################################

                 # 待开发  减少相似帧判断
                # if frame_number != 1:
                #     prev_frame = current_frame
                #     current_frame = NPcurrent_landmarks
                # else:
                    
                #     prev_frame = NPcurrent_landmarks
                #     current_frame = NPcurrent_landmarks

                # is_keyframe = keyframes(prev_frame, current_frame, 10)
                # print(is_keyframe)
                
                # print(NPcurrent_landmarks)     
                # result = set_compare(self.face,NPcurrent_landmarks)
                
                # if result is True:
                #     print(result)

                # if result
                  
                
####################################################################################################################################################

                return True
        return False
    

    # 绘制骨骼
    def draw_skeleton(self):
        # 定义骨架连接
        connections = [
            # 头部和躯干
            (0, 1), (1, 2), (2, 3), (3, 7),
            (0, 4), (4, 5), (5, 6), (6, 8),
            
            # 躯干
            (11, 12), (12, 14), (14, 16),  # 右臂
            (11, 13), (13, 15),  # 左臂
            (11, 23), (23, 25), (25, 27), (27, 29), (27, 31), (29,31),# 左腿
            (12, 24), (24, 26), (26, 28), (28, 30), (28, 32), (30,32), # 右腿
            
            # 身体中线
            (23, 24),  # 臀部
            (11, 12)   # 肩部
        ]
        
        # 绘制连接线
        for connection in connections:
            start = connection[0]
            end = connection[1]
            x_coords = [self.landmarks_3d[0][start], self.landmarks_3d[0][end]]
            y_coords = [self.landmarks_3d[1][start], self.landmarks_3d[1][end]]
            z_coords = [self.landmarks_3d[2][start], self.landmarks_3d[2][end]]
            self.ax.plot(x_coords, y_coords, z_coords, 'b-', linewidth=2)  # 'b-'：线段选择蓝色  linewidth=2：线宽=2

    def update_plot(self, frame):
        if self.process_frame():
            if self.landmarks_3d:
                self.ax.clear()
                
                # 绘制关键点
                self.ax.scatter(
                    self.landmarks_3d[0],
                    self.landmarks_3d[1],
                    self.landmarks_3d[2],
                    c='red',
                    marker='o'
                )
                
                # 绘制骨架
                self.draw_skeleton()
                
                # 设置坐标轴范围和标签
                self.ax.set_xlim([-1, 1])
                self.ax.set_ylim([-1, 1])
                self.ax.set_zlim([-1, 1])
                self.ax.set_xlabel('X')
                self.ax.set_ylabel('Z')
                self.ax.set_zlabel('Y')
                
                # 使用存储的视角参数
                self.ax.view_init(elev=self.elev, azim=self.azim) 
        else: 
            # 运行完后，运行__def__将csv关掉
            PoseDetector.__del__(self)
            print("完成姿态识别")


     
    def get_landmarks_3d(self):
        # 返回当前帧的关键点3D坐标
        return self.landmarks_3d
    

    def start_animation(self):
        anim = FuncAnimation(
            self.fig,
            self.update_plot,
            interval=1,   #每隔1ms更新一次图形
            blit=False
        )
        # plt.show()
        try:
            anim.save('animation.mp4', writer='ffmpeg', fps=30)
        except Exception as e:
            # print(type(aaa))
            print(f"Error saving animation: {e}")
            print(777)
            return None


    def __del__(self):
        self.cap.release()
        cv2.destroyAllWindows()
        # 关闭三个csv文件
        try:
            self.pose.close()
            # self.csv_AllLandmarks.close()
            # self.csv_NeedData.close()
            # self.csv_AttackData.close()
            self.csv[0].close()
            self.csv[1].close()
            self.csv[2].close()


        except Exception:
            self.pose.close()
            # self.csv_AllLandmarks.close()
            # self.csv_NeedData.close()
            # self.csv_AttackData.close()
            self.csv[0].close()
            self.csv[1].close()
            self.csv[2].close()


# 使用示例
def main():

    height_real = int(input("请输入真实身高"))
    side = int(input("选择需要识别的人(左边的人=1，右边的人=2，单人视频=0)"))
    video_path = "video/Pixel_Low.mp4"  # 替换为你的视频路径
    detector = PoseDetector(video_path,height_real,side)
    # 启动动画
    detector.start_animation()
    

if __name__ == "__main__":
    main()